﻿using Microsoft.AspNet.SignalR;
using Microsoft.AspNet.SignalR.Hubs;
using Microsoft.AspNet.SignalR.Owin;
using Microsoft.Owin.Security;
using Microsoft.Owin.Security.DataHandler;
using Microsoft.Owin.Security.DataProtection;
using Microsoft.Owin.Security.OAuth;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Security.Claims;
using System.Web;

namespace YlisChat.App_Start
{
    public class MachineKeyProtector : IDataProtector
    {
        private readonly string[] _purpose =
        {
            typeof(OAuthAuthorizationServerMiddleware).Namespace,
            "Access_Token",
            "v1"
        };

        public byte[] Protect(byte[] userData)
        {
            throw new NotImplementedException();
        }

        public byte[] Unprotect(byte[] protectedData)
        {
            return System.Web.Security.MachineKey.Unprotect(protectedData, _purpose);
        }
    }

    // Adopted from http://geekswithblogs.net/shaunxu/archive/2014/05/27/set-context-user-principal-for-customized-authentication-in-signalr.aspx
    // and from http://stackoverflow.com/questions/20585872/get-iprincipal-from-oauth-bearer-token-in-owin
    public class CookieBearerAuthorizeAttribute : AuthorizeAttribute
    {
        public override bool AuthorizeHubConnection(HubDescriptor hubDescriptor, IRequest request)
        {
            //DpapiDataProtectionProvider dataProtectionProvider = new DpapiDataProtectionProvider();
            //TicketDataFormat secureDataFormat = new TicketDataFormat(dataProtectionProvider.Create());
            TicketDataFormat secureDataFormat = new TicketDataFormat(new MachineKeyProtector());
            // authenticate by using bearer token in cookies
            Cookie bearer;
            if (!request.Cookies.TryGetValue("Bearer", out bearer))
                return false;
            
            //var token = request.QueryString.Get(WebApiConfig.AuthenticationType);
            AuthenticationTicket ticket = secureDataFormat.Unprotect(bearer.Value);
            if (ticket != null && ticket.Identity != null && ticket.Identity.IsAuthenticated)
            {
                // set the authenticated user principal into environment so that it can be used in the future
                request.Environment["server.User"] = new ClaimsPrincipal(ticket.Identity);
                return true;
            }
            else
            {
                return false;
            }
        }

        public override bool AuthorizeHubMethodInvocation(IHubIncomingInvokerContext hubIncomingInvokerContext, bool appliesToMethod)
        {
            string connectionId = hubIncomingInvokerContext.Hub.Context.ConnectionId;
            // check the authenticated user principal from environment
            IDictionary<string, object> environment = hubIncomingInvokerContext.Hub.Context.Request.Environment;
            ClaimsPrincipal principal = environment["server.User"] as ClaimsPrincipal;
            if (principal != null && principal.Identity != null && principal.Identity.IsAuthenticated)
            {
                // create a new HubCallerContext instance with the principal generated from token
                // and replace the current context so that in hubs we can retrieve current user identity
                hubIncomingInvokerContext.Hub.Context = new HubCallerContext(new ServerRequest(environment), connectionId);
                return true;
            }
            else
            {
                return false;
            }
        }
    }
}